from paddlenlp.transformers import UIE, AutoTokenizer
import paddle

# 加载动态图模型和分词器
from paddlenlp.transformers import UIE, AutoTokenizer
import paddle

# 加载动态图模型和分词器
model = UIE.from_pretrained("D:\\paddle_uie")  # 替换为你的本地模型路径
tokenizer = AutoTokenizer.from_pretrained("D:\\paddle_uie")

# 定义输入规格
input_spec = [
    paddle.static.InputSpec(shape=[None, None], dtype="int64", name="input_ids"),
    paddle.static.InputSpec(shape=[None, None], dtype="int64", name="token_type_ids")
]

# 导出静态图模型
paddle.jit.save(model, "D:\\paddle_uie\\static\\inference", input_spec=input_spec)

from paddlenlp import Taskflow

# 加载本地模型
ie = Taskflow(
    task="information_extraction",
    schema=["时间", "地点", "人物"],
    task_path=r"D:\paddle_uie"  # 替换为你的本地模型路径
)

# 测试推理
result = ie("2024年北京冬奥会在鸟巢举行。")
print(result)